package com.shujia.mllib

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo5ImageTrain {
  def main(args: Array[String]): Unit = {

    val spark: SparkSession = SparkSession
      .builder()
      .master("local[8]")
      .appName("image")
      .getOrCreate()

    import spark.implicits._
    import org.apache.spark.sql.functions._


    /**
      * 加载数据
      *
      */

    val data: DataFrame = spark
      .read
      .format("libsvm")
      .load("Spark/data/image_data")


    /**
      * 切分训练集和测试集
      *
      */

    val split: Array[Dataset[Row]] = data.randomSplit(Array(0.7, 0.3))
    val train: Dataset[Row] = split(0)
    val test: Dataset[Row] = split(1)


    /**
      * 选择算法
      *
      */

    val logisticRegression = new LogisticRegression()


    //将训练集带入算法训练模型
    val model: LogisticRegressionModel = logisticRegression.fit(train)


    //使用模型对测试集进行测试
    val testDF: DataFrame = model.transform(test)


    //计算准确率
    testDF
      .select(sum(when($"label" === $"prediction", 1.0).otherwise(0.0)) / count($"label"))
      .show()


    //保存模型
    model.write
      .overwrite()
      .save("Spark/data/image_model")


  }

}
